module fp16_pe (
    input  logic      clk,           // 时钟信号
    input  logic      rst_n,         // 复位信号（低电平有效）
    input  logic      valid_in,      // 输入数据有效
    input  pack_16B_t vector_a,      // 8个FP16数据A: a[7], a[6], ..., a[0]
    input  pack_16B_t vector_b,      // 8个FP16数据B: b[7], b[6], ..., b[0]
    input  fp32_t     bias_c,        // FP32偏置值C
    output logic      valid_out,     // 输出数据有效
    output fp32_t     result,        // FP32最终结果
    output logic      exception,     // 异常标志
    // 新增：将输入的向量传递给下一级
    output pack_16B_t vector_a_out,
    output pack_16B_t vector_b_out
);

    // 内部信号定义
    logic [15:0] fp16_mul_results[8];  // 8个FP16乘法结果 (a[i] * b[i])
    logic [7:0] mul_exceptions;  // 8个乘法器的异常标志位
    logic [7:0] mul_overflows;  // 8个乘法器的溢出标志位
    logic [7:0] mul_underflows;  // 8个乘法器的下溢标志位
    logic [7:0] mul_valid_out;  // 8个乘法器的输出有效信号

    logic [127:0] acc_operands;  // 累加器操作数输入 (8个FP16结果打包)
    logic [31:0] fp32_acc_result;  // 累加器直接输出FP32结果
    logic acc_exception;  // 累加器异常标志
    logic acc_valid_out;  // 累加器输出有效信号

    logic final_add_exception;  // 最终FP32加法异常标志
    logic final_add_valid_out;  // 最终FP32加法输出有效信号

    // 中间级的有效信号，用于控制流水线
    logic mul_stage_valid;  // 乘法阶段输出有效
    logic acc_stage_valid;  // 累加阶段输出有效

    // 1. 实例化8个FP16乘法器 - 计算向量点积的乘法部分
    generate
        for (genvar i = 0; i < 8; i++) begin : gen_fp16_multipliers
            fp16_mul u_fp16_mul (
                .clk      (clk),
                .rst_n    (rst_n),
                .valid_in (valid_in),             // 输入有效信号
                .a_operand(vector_a.fp16[i]),     // 提取向量A的第i个FP16元素
                .b_operand(vector_b.fp16[i]),     // 提取向量B的第i个FP16元素
                .result   (fp16_mul_results[i]),  // 第i个乘法结果 (FP16)
                .Exception(mul_exceptions[i]),    // 第i个乘法器异常标志
                .Overflow (mul_overflows[i]),     // 第i个乘法器溢出标志
                .Underflow(mul_underflows[i]),    // 第i个乘法器下溢标志
                .valid_out(mul_valid_out[i])      // 第i个乘法器输出有效信号
            );
        end
    endgenerate

    // 乘法阶段的有效信号 - 所有乘法器都完成时才有效
    always_comb begin
        mul_stage_valid = &mul_valid_out;  // 所有乘法器输出都有效时
    end

    // 2. 将8个FP16乘法结果打包为累加器输入
    always_comb begin
        for (int i = 0; i < 8; i++) begin
            acc_operands[i*16+:16] = fp16_mul_results[i];
        end
    end

    // 3. 实例化FP16累加器 - 将8个FP16乘法结果累加并转换为FP32
    fp16_acc #(
        .INPUTS(8),
        .PIPELINE_STAGES(4)
    ) u_fp16_acc (
        .clk      (clk),
        .rst_n    (rst_n),
        .valid_in (mul_stage_valid),  // 使用乘法阶段的有效信号
        .operands (acc_operands),     // 8个FP16乘法结果作为操作数
        .valid_out(acc_valid_out),    // 累加器输出有效信号
        .result   (fp32_acc_result),  // 累加结果直接输出为FP32格式
        .exception(acc_exception)     // 累加过程中的异常标志
    );

    // 4. 最终FP32加法 - 将累加结果与偏置值相加
    fp32_add u_fp32_add (
        .clk      (clk),
        .rst_n    (rst_n),
        .valid_in (acc_valid_out),        // 使用累加器的输出有效信号
        .a_operand(fp32_acc_result),      // 累加器输出的FP32结果
        .b_operand(bias_c),               // FP32偏置值
        .valid_out(final_add_valid_out),  // 最终加法输出有效信号
        .result   (result),               // 最终输出结果 (FP32)
        .exception(final_add_exception)   // 最终加法异常标志
    );

    // 5. 输出控制
    assign valid_out = final_add_valid_out;

    // 6. 异常信号汇总 - 只在输出有效时汇总异常
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            exception <= 1'b0;
        end else if (final_add_valid_out) begin
            exception <= (|mul_exceptions) |  // 8个乘法器中任一异常
            acc_exception |  // 累加器异常
            final_add_exception;  // 最终加法异常
        end
    end

    // 在模块末尾添加：延迟一个周期传递输入向量
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            vector_a_out <= '0;
            vector_b_out <= '0;
        end else begin
            vector_a_out <= vector_a;
            vector_b_out <= vector_b;
        end
    end

endmodule
